clip_param_keys = {
    'positional_embedding',
    'text_projection',
    'visual.class_embedding',
    'visual.positional_embedding',
    'visual.conv1.weight',
    'visual.ln_pre.weight',
    'visual.ln_pre.bias',
    'visual.transformer.resblocks.0.attn.in_proj_weight',
    'visual.transformer.resblocks.0.attn.in_proj_bias',
    'visual.transformer.resblocks.0.attn.out_proj.weight',
    'visual.transformer.resblocks.0.attn.out_proj.bias',
    'visual.transformer.resblocks.0.ln_1.weight',
    'visual.transformer.resblocks.0.ln_1.bias',
    'visual.transformer.resblocks.0.mlp.c_fc.weight',
    'visual.transformer.resblocks.0.mlp.c_fc.bias',
    'visual.transformer.resblocks.0.mlp.c_proj.weight',
    'visual.transformer.resblocks.0.mlp.c_proj.bias',
    'visual.transformer.resblocks.0.ln_2.weight',
    'visual.transformer.resblocks.0.ln_2.bias',
    'visual.transformer.resblocks.1.attn.in_proj_weight',
    'visual.transformer.resblocks.1.attn.in_proj_bias',
    'visual.transformer.resblocks.1.attn.out_proj.weight',
    'visual.transformer.resblocks.1.attn.out_proj.bias',
    'visual.transformer.resblocks.1.ln_1.weight',
    'visual.transformer.resblocks.1.ln_1.bias',
    'visual.transformer.resblocks.1.mlp.c_fc.weight',
    'visual.transformer.resblocks.1.mlp.c_fc.bias',
    'visual.transformer.resblocks.1.mlp.c_proj.weight',
    'visual.transformer.resblocks.1.mlp.c_proj.bias',
    'visual.transformer.resblocks.1.ln_2.weight',
    'visual.transformer.resblocks.1.ln_2.bias',
    'visual.transformer.resblocks.2.attn.in_proj_weight',
    'visual.transformer.resblocks.2.attn.in_proj_bias',
    'visual.transformer.resblocks.2.attn.out_proj.weight',
    'visual.transformer.resblocks.2.attn.out_proj.bias',
    'visual.transformer.resblocks.2.ln_1.weight',
    'visual.transformer.resblocks.2.ln_1.bias',
    'visual.transformer.resblocks.2.mlp.c_fc.weight',
    'visual.transformer.resblocks.2.mlp.c_fc.bias',
    'visual.transformer.resblocks.2.mlp.c_proj.weight',
    'visual.transformer.resblocks.2.mlp.c_proj.bias',
    'visual.transformer.resblocks.2.ln_2.weight',
    'visual.transformer.resblocks.2.ln_2.bias',
    'visual.transformer.resblocks.3.attn.in_proj_weight',
    'visual.transformer.resblocks.3.attn.in_proj_bias',
    'visual.transformer.resblocks.3.attn.out_proj.weight',
    'visual.transformer.resblocks.3.attn.out_proj.bias',
    'visual.transformer.resblocks.3.ln_1.weight',
    'visual.transformer.resblocks.3.ln_1.bias',
    'visual.transformer.resblocks.3.mlp.c_fc.weight',
    'visual.transformer.resblocks.3.mlp.c_fc.bias',
    'visual.transformer.resblocks.3.mlp.c_proj.weight',
    'visual.transformer.resblocks.3.mlp.c_proj.bias',
    'visual.transformer.resblocks.3.ln_2.weight',
    'visual.transformer.resblocks.3.ln_2.bias',
    'visual.transformer.resblocks.4.attn.in_proj_weight',
    'visual.transformer.resblocks.4.attn.in_proj_bias',
    'visual.transformer.resblocks.4.attn.out_proj.weight',
    'visual.transformer.resblocks.4.attn.out_proj.bias',
    'visual.transformer.resblocks.4.ln_1.weight',
    'visual.transformer.resblocks.4.ln_1.bias',
    'visual.transformer.resblocks.4.mlp.c_fc.weight',
    'visual.transformer.resblocks.4.mlp.c_fc.bias',
    'visual.transformer.resblocks.4.mlp.c_proj.weight',
    'visual.transformer.resblocks.4.mlp.c_proj.bias',
    'visual.transformer.resblocks.4.ln_2.weight',
    'visual.transformer.resblocks.4.ln_2.bias',
    'visual.transformer.resblocks.5.attn.in_proj_weight',
    'visual.transformer.resblocks.5.attn.in_proj_bias',
    'visual.transformer.resblocks.5.attn.out_proj.weight',
    'visual.transformer.resblocks.5.attn.out_proj.bias',
    'visual.transformer.resblocks.5.ln_1.weight',
    'visual.transformer.resblocks.5.ln_1.bias',
    'visual.transformer.resblocks.5.mlp.c_fc.weight',
    'visual.transformer.resblocks.5.mlp.c_fc.bias',
    'visual.transformer.resblocks.5.mlp.c_proj.weight',
    'visual.transformer.resblocks.5.mlp.c_proj.bias',
    'visual.transformer.resblocks.5.ln_2.weight',
    'visual.transformer.resblocks.5.ln_2.bias',
    'visual.transformer.resblocks.6.attn.in_proj_weight',
    'visual.transformer.resblocks.6.attn.in_proj_bias',
    'visual.transformer.resblocks.6.attn.out_proj.weight',
    'visual.transformer.resblocks.6.attn.out_proj.bias',
    'visual.transformer.resblocks.6.ln_1.weight',
    'visual.transformer.resblocks.6.ln_1.bias',
    'visual.transformer.resblocks.6.mlp.c_fc.weight',
    'visual.transformer.resblocks.6.mlp.c_fc.bias',
    'visual.transformer.resblocks.6.mlp.c_proj.weight',
    'visual.transformer.resblocks.6.mlp.c_proj.bias',
    'visual.transformer.resblocks.6.ln_2.weight',
    'visual.transformer.resblocks.6.ln_2.bias',
    'visual.transformer.resblocks.7.attn.in_proj_weight',
    'visual.transformer.resblocks.7.attn.in_proj_bias',
    'visual.transformer.resblocks.7.attn.out_proj.weight',
    'visual.transformer.resblocks.7.attn.out_proj.bias',
    'visual.transformer.resblocks.7.ln_1.weight',
    'visual.transformer.resblocks.7.ln_1.bias',
    'visual.transformer.resblocks.7.mlp.c_fc.weight',
    'visual.transformer.resblocks.7.mlp.c_fc.bias',
    'visual.transformer.resblocks.7.mlp.c_proj.weight',
    'visual.transformer.resblocks.7.mlp.c_proj.bias',
    'visual.transformer.resblocks.7.ln_2.weight',
    'visual.transformer.resblocks.7.ln_2.bias',
    'visual.transformer.resblocks.8.attn.in_proj_weight',
    'visual.transformer.resblocks.8.attn.in_proj_bias',
    'visual.transformer.resblocks.8.attn.out_proj.weight',
    'visual.transformer.resblocks.8.attn.out_proj.bias',
    'visual.transformer.resblocks.8.ln_1.weight',
    'visual.transformer.resblocks.8.ln_1.bias',
    'visual.transformer.resblocks.8.mlp.c_fc.weight',
    'visual.transformer.resblocks.8.mlp.c_fc.bias',
    'visual.transformer.resblocks.8.mlp.c_proj.weight',
    'visual.transformer.resblocks.8.mlp.c_proj.bias',
    'visual.transformer.resblocks.8.ln_2.weight',
    'visual.transformer.resblocks.8.ln_2.bias',
    'visual.transformer.resblocks.9.attn.in_proj_weight',
    'visual.transformer.resblocks.9.attn.in_proj_bias',
    'visual.transformer.resblocks.9.attn.out_proj.weight',
    'visual.transformer.resblocks.9.attn.out_proj.bias',
    'visual.transformer.resblocks.9.ln_1.weight',
    'visual.transformer.resblocks.9.ln_1.bias',
    'visual.transformer.resblocks.9.mlp.c_fc.weight',
    'visual.transformer.resblocks.9.mlp.c_fc.bias',
    'visual.transformer.resblocks.9.mlp.c_proj.weight',
    'visual.transformer.resblocks.9.mlp.c_proj.bias',
    'visual.transformer.resblocks.9.ln_2.weight',
    'visual.transformer.resblocks.9.ln_2.bias',
    'visual.transformer.resblocks.10.attn.in_proj_weight',
    'visual.transformer.resblocks.10.attn.in_proj_bias',
    'visual.transformer.resblocks.10.attn.out_proj.weight',
    'visual.transformer.resblocks.10.attn.out_proj.bias',
    'visual.transformer.resblocks.10.ln_1.weight',
    'visual.transformer.resblocks.10.ln_1.bias',
    'visual.transformer.resblocks.10.mlp.c_fc.weight',
    'visual.transformer.resblocks.10.mlp.c_fc.bias',
    'visual.transformer.resblocks.10.mlp.c_proj.weight',
    'visual.transformer.resblocks.10.mlp.c_proj.bias',
    'visual.transformer.resblocks.10.ln_2.weight',
    'visual.transformer.resblocks.10.ln_2.bias',
    'visual.transformer.resblocks.11.attn.in_proj_weight',
    'visual.transformer.resblocks.11.attn.in_proj_bias',
    'visual.transformer.resblocks.11.attn.out_proj.weight',
    'visual.transformer.resblocks.11.attn.out_proj.bias',
    'visual.transformer.resblocks.11.ln_1.weight',
    'visual.transformer.resblocks.11.ln_1.bias',
    'visual.transformer.resblocks.11.mlp.c_fc.weight',
    'visual.transformer.resblocks.11.mlp.c_fc.bias',
    'visual.transformer.resblocks.11.mlp.c_proj.weight',
    'visual.transformer.resblocks.11.mlp.c_proj.bias',
    'visual.transformer.resblocks.11.ln_2.weight',
    'visual.transformer.resblocks.11.ln_2.bias',
    'visual.ln_post.weight',
    'visual.ln_post.bias',
    'visual_ln_post.weight',
    'visual_ln_post.bias',
    'transformer.resblocks.0.attn.in_proj_weight',
    'transformer.resblocks.0.attn.in_proj_bias',
    'transformer.resblocks.0.attn.out_proj.weight',
    'transformer.resblocks.0.attn.out_proj.bias',
    'transformer.resblocks.0.ln_1.weight',
    'transformer.resblocks.0.ln_1.bias',
    'transformer.resblocks.0.mlp.c_fc.weight',
    'transformer.resblocks.0.mlp.c_fc.bias',
    'transformer.resblocks.0.mlp.c_proj.weight',
    'transformer.resblocks.0.mlp.c_proj.bias',
    'transformer.resblocks.0.ln_2.weight',
    'transformer.resblocks.0.ln_2.bias',
    'transformer.resblocks.1.attn.in_proj_weight',
    'transformer.resblocks.1.attn.in_proj_bias',
    'transformer.resblocks.1.attn.out_proj.weight',
    'transformer.resblocks.1.attn.out_proj.bias',
    'transformer.resblocks.1.ln_1.weight',
    'transformer.resblocks.1.ln_1.bias',
    'transformer.resblocks.1.mlp.c_fc.weight',
    'transformer.resblocks.1.mlp.c_fc.bias',
    'transformer.resblocks.1.mlp.c_proj.weight',
    'transformer.resblocks.1.mlp.c_proj.bias',
    'transformer.resblocks.1.ln_2.weight',
    'transformer.resblocks.1.ln_2.bias',
    'transformer.resblocks.2.attn.in_proj_weight',
    'transformer.resblocks.2.attn.in_proj_bias',
    'transformer.resblocks.2.attn.out_proj.weight',
    'transformer.resblocks.2.attn.out_proj.bias',
    'transformer.resblocks.2.ln_1.weight',
    'transformer.resblocks.2.ln_1.bias',
    'transformer.resblocks.2.mlp.c_fc.weight',
    'transformer.resblocks.2.mlp.c_fc.bias',
    'transformer.resblocks.2.mlp.c_proj.weight',
    'transformer.resblocks.2.mlp.c_proj.bias',
    'transformer.resblocks.2.ln_2.weight',
    'transformer.resblocks.2.ln_2.bias',
    'transformer.resblocks.3.attn.in_proj_weight',
    'transformer.resblocks.3.attn.in_proj_bias',
    'transformer.resblocks.3.attn.out_proj.weight',
    'transformer.resblocks.3.attn.out_proj.bias',
    'transformer.resblocks.3.ln_1.weight',
    'transformer.resblocks.3.ln_1.bias',
    'transformer.resblocks.3.mlp.c_fc.weight',
    'transformer.resblocks.3.mlp.c_fc.bias',
    'transformer.resblocks.3.mlp.c_proj.weight',
    'transformer.resblocks.3.mlp.c_proj.bias',
    'transformer.resblocks.3.ln_2.weight',
    'transformer.resblocks.3.ln_2.bias',
    'transformer.resblocks.4.attn.in_proj_weight',
    'transformer.resblocks.4.attn.in_proj_bias',
    'transformer.resblocks.4.attn.out_proj.weight',
    'transformer.resblocks.4.attn.out_proj.bias',
    'transformer.resblocks.4.ln_1.weight',
    'transformer.resblocks.4.ln_1.bias',
    'transformer.resblocks.4.mlp.c_fc.weight',
    'transformer.resblocks.4.mlp.c_fc.bias',
    'transformer.resblocks.4.mlp.c_proj.weight',
    'transformer.resblocks.4.mlp.c_proj.bias',
    'transformer.resblocks.4.ln_2.weight',
    'transformer.resblocks.4.ln_2.bias',
    'transformer.resblocks.5.attn.in_proj_weight',
    'transformer.resblocks.5.attn.in_proj_bias',
    'transformer.resblocks.5.attn.out_proj.weight',
    'transformer.resblocks.5.attn.out_proj.bias',
    'transformer.resblocks.5.ln_1.weight',
    'transformer.resblocks.5.ln_1.bias',
    'transformer.resblocks.5.mlp.c_fc.weight',
    'transformer.resblocks.5.mlp.c_fc.bias',
    'transformer.resblocks.5.mlp.c_proj.weight',
    'transformer.resblocks.5.mlp.c_proj.bias',
    'transformer.resblocks.5.ln_2.weight',
    'transformer.resblocks.5.ln_2.bias',
    'transformer.resblocks.6.attn.in_proj_weight',
    'transformer.resblocks.6.attn.in_proj_bias',
    'transformer.resblocks.6.attn.out_proj.weight',
    'transformer.resblocks.6.attn.out_proj.bias',
    'transformer.resblocks.6.ln_1.weight',
    'transformer.resblocks.6.ln_1.bias',
    'transformer.resblocks.6.mlp.c_fc.weight',
    'transformer.resblocks.6.mlp.c_fc.bias',
    'transformer.resblocks.6.mlp.c_proj.weight',
    'transformer.resblocks.6.mlp.c_proj.bias',
    'transformer.resblocks.6.ln_2.weight',
    'transformer.resblocks.6.ln_2.bias',
    'transformer.resblocks.7.attn.in_proj_weight',
    'transformer.resblocks.7.attn.in_proj_bias',
    'transformer.resblocks.7.attn.out_proj.weight',
    'transformer.resblocks.7.attn.out_proj.bias',
    'transformer.resblocks.7.ln_1.weight',
    'transformer.resblocks.7.ln_1.bias',
    'transformer.resblocks.7.mlp.c_fc.weight',
    'transformer.resblocks.7.mlp.c_fc.bias',
    'transformer.resblocks.7.mlp.c_proj.weight',
    'transformer.resblocks.7.mlp.c_proj.bias',
    'transformer.resblocks.7.ln_2.weight',
    'transformer.resblocks.7.ln_2.bias',
    'transformer.resblocks.8.attn.in_proj_weight',
    'transformer.resblocks.8.attn.in_proj_bias',
    'transformer.resblocks.8.attn.out_proj.weight',
    'transformer.resblocks.8.attn.out_proj.bias',
    'transformer.resblocks.8.ln_1.weight',
    'transformer.resblocks.8.ln_1.bias',
    'transformer.resblocks.8.mlp.c_fc.weight',
    'transformer.resblocks.8.mlp.c_fc.bias',
    'transformer.resblocks.8.mlp.c_proj.weight',
    'transformer.resblocks.8.mlp.c_proj.bias',
    'transformer.resblocks.8.ln_2.weight',
    'transformer.resblocks.8.ln_2.bias',
    'transformer.resblocks.9.attn.in_proj_weight',
    'transformer.resblocks.9.attn.in_proj_bias',
    'transformer.resblocks.9.attn.out_proj.weight',
    'transformer.resblocks.9.attn.out_proj.bias',
    'transformer.resblocks.9.ln_1.weight',
    'transformer.resblocks.9.ln_1.bias',
    'transformer.resblocks.9.mlp.c_fc.weight',
    'transformer.resblocks.9.mlp.c_fc.bias',
    'transformer.resblocks.9.mlp.c_proj.weight',
    'transformer.resblocks.9.mlp.c_proj.bias',
    'transformer.resblocks.9.ln_2.weight',
    'transformer.resblocks.9.ln_2.bias',
    'transformer.resblocks.10.attn.in_proj_weight',
    'transformer.resblocks.10.attn.in_proj_bias',
    'transformer.resblocks.10.attn.out_proj.weight',
    'transformer.resblocks.10.attn.out_proj.bias',
    'transformer.resblocks.10.ln_1.weight',
    'transformer.resblocks.10.ln_1.bias',
    'transformer.resblocks.10.mlp.c_fc.weight',
    'transformer.resblocks.10.mlp.c_fc.bias',
    'transformer.resblocks.10.mlp.c_proj.weight',
    'transformer.resblocks.10.mlp.c_proj.bias',
    'transformer.resblocks.10.ln_2.weight',
    'transformer.resblocks.10.ln_2.bias',
    'transformer.resblocks.11.attn.in_proj_weight',
    'transformer.resblocks.11.attn.in_proj_bias',
    'transformer.resblocks.11.attn.out_proj.weight',
    'transformer.resblocks.11.attn.out_proj.bias',
    'transformer.resblocks.11.ln_1.weight',
    'transformer.resblocks.11.ln_1.bias',
    'transformer.resblocks.11.mlp.c_fc.weight',
    'transformer.resblocks.11.mlp.c_fc.bias',
    'transformer.resblocks.11.mlp.c_proj.weight',
    'transformer.resblocks.11.mlp.c_proj.bias',
    'transformer.resblocks.11.ln_2.weight',
    'transformer.resblocks.11.ln_2.bias',
    'token_embedding.weight',
    'visual.proj',
    'visual_proj',
    'ln_final.weight',
    'ln_final.bias',
}


def gradually_freeze_by_layer(model, global_step, interval, unfreeze_type="both"):
    try:
        visual_n_layers = model.clip.visual.transformer.layers
    except:
        visual_n_layers = model.visual.transformer.layers

    try:
        text_n_layers = model.clip.transformer.layers
    except:
        text_n_layers = model.transformer.layers

    if interval <= 0:
        return

    layer_to_unfreeze = global_step // interval
    if (
        visual_n_layers + 1 < layer_to_unfreeze
        and text_n_layers + 1 < layer_to_unfreeze
    ):  # Do nothing
        return
    embeddings = [
        "token_embedding",
        "visual.positional_embedding" "visual.class_embedding",
        "positional_embedding",
    ]
    visual_first_keys = [
        "visual.conv1.weight",
        "visual.ln_pre",
    ]
    text_first_keys = []
    visual_last_keys = ["visual.ln_post", "visual_ln_post" "visual_proj", "visual.proj"]
    text_last_keys = ["ln_final"]

    def set_unfreeze(n, p, l):
        if any(x in n for x in l):
            p.requires_grad = True

    for n, p in model.named_parameters():
        if layer_to_unfreeze == 0:  # At the beginning of training
            set_unfreeze(n, p, embeddings)  # Unfreeze embeddings
            if unfreeze_type in ["both", "visual"]:
                set_unfreeze(n, p, visual_last_keys)
            if unfreeze_type in ["both", "text"]:  # Unfreeze text last layers
                set_unfreeze(n, p, text_last_keys)
            continue

        if unfreeze_type in ["both", "visual"]:
            if layer_to_unfreeze <= visual_n_layers:  # Unfreeze n-th visual layer
                param_prefix = f"visual.transformer.resblocks.{visual_n_layers - layer_to_unfreeze}"
                set_unfreeze(n, p, [param_prefix])
            else:  # Unfreeze first visual layers
                set_unfreeze(n, p, visual_first_keys)
        if unfreeze_type in ["both", "text"]:
            if layer_to_unfreeze <= text_n_layers:  # Unfreeze n-th text layer
                param_prefix = (
                    f"transformer.resblocks.{visual_n_layers - layer_to_unfreeze}"
                )
                set_unfreeze(n, p, [param_prefix])
            else:  # Unfreeze first visual layers
                set_unfreeze(n, p, text_first_keys)
